import numpy as np
import cv2 as cv
import random
import torch
import os
from myNet import *
drawing = False # true if mouse is pressed
clear = False # set true to clear image
# mouse callback function
def draw_circle(event,x,y,flags,param):
    global ix,iy,drawing,img,clear
    if event == cv.EVENT_LBUTTONDOWN:
        drawing = True
        if clear:
            img = np.zeros((256,256), np.uint8)
            clear = False
    elif event == cv.EVENT_MOUSEMOVE:
        if drawing == True:
            cv.circle(img,(x,y),15,(255,),-1)
    elif event == cv.EVENT_LBUTTONUP:
        drawing = False
    elif event == cv.EVENT_LBUTTONDBLCLK:
        parse_input_image(img)
        clear = True

# TODO：
# 以下代码是最近邻方法（可以正常运行），请使用上一步训练得到的全连接网络模型替换下列代码，对比下两者的不同
device = torch.device('cpu')
model = Net().to(device)#实例化网络
model.load_state_dict(torch.load('mnist_cnn.pt'))#将训练出来的参数加入model里面
model.eval()
# training_data, training_target = torch.load(r'C:\Users\lenovo1\Desktop\training.pt')
# training_data = training_data.reshape((60000, -1)) / 255.0
def parse_input_image(img):   # 这里不要改动
    # 注释：绘制的图像大小是256*256，需要缩放到跟mnist数据集一样的28*28；并把数据从[0~255]调整为[0~1.0]
    img_resize = cv.resize(img, (28, 28)).reshape(-1) / 255.0 # 注意这里是numpy的array
    img_resize = img_resize.astype(np.float32)
    img_resize = torch.from_numpy(img_resize)
    data = img_resize.reshape((1,1,28,28))#变为1，1，28，28的张量
    # 开始预测
    output = model(data)
    predict_result = output.argmax(dim=1, keepdim=True)
    print(f'Your enter a {predict_result.item()}')
# TODO结束  

tmp_file_num = 1
def save_image_to_folder(img, folder_name):
    global tmp_file_num
    filename = os.path.join(folder_name, f"{tmp_file_num}.png")
    while os.path.exists(filename):
        tmp_file_num += 1
        filename = os.path.join(folder_name, f"{tmp_file_num}.png")
    cv.imwrite(filename, img)
    print(f"save image to {filename}")
    tmp_file_num += 1

img = np.zeros((256,256), np.uint8)
cv.namedWindow('image')
cv.setMouseCallback('image',draw_circle)
while(1):
    cv.imshow('image',img)
    k = cv.waitKey(20) & 0xFF  
    if k == 27: # press Esc to exit
        break
    elif k == 32 or k == 13:  # press blank or enter to complete drawing
        parse_input_image(img)
        clear = True
    elif k == 83 or k == 115:
        save_image_to_folder(img, "mydata")
        clear = True
cv.destroyAllWindows()

# opencv code from: https://docs.opencv.org/4.x/db/d5b/tutorial_py_mouse_handling.html